import torch
import onnx
import numpy as np
import onnxruntime as rt

from args import args
from build_net import make_model


def cosine_similarity(vector1, vector2):
    vector1 = vector1.flatten()
    vector2 = vector2.flatten()

    unit_vector1 = vector1 / np.linalg.norm(vector1)
    unit_vector2 = vector2 / np.linalg.norm(vector2)

    similarity = np.dot(unit_vector1, unit_vector2)
    return similarity


if __name__ == '__main__':
    model = make_model(args)
    model.load_state_dict(torch.load(args.model_path, map_location=torch.device('cpu')))
    model.eval
    
    # 测试模型功能
    img = torch.zeros((1, 3, 224, 224))
    pred = model(img)
    
    # 转onnx
    inputs = (img)

    input_names = ['image']
    output_names = ['pred']

    # dynamic_axes = {'latents': {0: '-1'}}

    save_onnx = 'terrains1_bs1_512.onnx'

    torch.onnx.export(
        model,
        inputs,
        save_onnx,
        input_names=input_names,
        # dynamic_axes=dynamic_axes,
        output_names=output_names,
        opset_version=16,
    )
    
    print(f"[INFO] saved onnx: {save_onnx}")
    
    # 验证onnx精度
    sess = rt.InferenceSession(save_onnx, providers=['CPUExecutionProvider'])

    net_inputs = {'image': img.numpy()}

    net_output = sess.run(None, net_inputs)[0]

    print(f"[INFO] check onnx cosine_similarity: {cosine_similarity(pred[0].detach().numpy(), net_output)}")
    
    print(f'[INFO] refer cmd to om: atc --framework=5 --model={save_onnx} --output={save_onnx.split(".")[0]} --input_format=NCHW --input_shape="image:1,3,224,224" --log=error --soc_version=Ascend910B3')
    
    